import os
import xml.etree.ElementTree as ET
from PIL import Image


image_folder_path = r'D:\zhenghe_biaozhipai\JPEGImages'


annotation_folder_path = r'D:\zhenghe_biaozhipai\Annotations'

# 阈值
thresholds = 100


output_folder_path = r'D:\zhenghe_biaozhipai\crop_out\shaixuan'
os.makedirs(output_folder_path, exist_ok=True)

# 小于阈值数组保存的文件夹
rejected_folder_path = r'D:\zhenghe_biaozhipai\crop_out\fuhe'
os.makedirs(rejected_folder_path, exist_ok=True)


for filename in os.listdir(annotation_folder_path):
    if not filename.endswith('.xml'):
        continue

    xml_path = os.path.join(annotation_folder_path, filename)

    
    tree = ET.parse(xml_path)
    root = tree.getroot()

    # 获取图片文件名
    image_filename = root.find('filename').text
    image_path = os.path.join(image_folder_path, image_filename)

    
    image = Image.open(image_path)

   
    for object_element in root.findall('object'):
        # 类别信息
        class_name = object_element.find('name').text

        # 边界框信息
        bbox_element = object_element.find('bndbox')
        xmin = int(bbox_element.find('xmin').text)
        ymin = int(bbox_element.find('ymin').text)
        xmax = int(bbox_element.find('xmax').text)
        ymax = int(bbox_element.find('ymax').text)

        # 目标尺寸大小
        width = xmax - xmin
        height = ymax - ymin

        # 裁剪目标区域
        cropped_image = image.crop((xmin, ymin, xmax, ymax))

        
       
        if width >= thresholds and height >= thresholds:
            
            class_folder_path = os.path.join(output_folder_path, class_name, str(thresholds))
            os.makedirs(class_folder_path, exist_ok=True)

            #保存路径
            crop_filename = f"{image_filename.split('.')[0]}_{xmin}_{ymin}_{xmax}_{ymax}.jpg"
            crop_path = os.path.join(class_folder_path, crop_filename)

            
            cropped_image.save(crop_path)

            print(f"Cropped image saved: {crop_filename}")
            
        else:
            # 不在阈值数组内的图片
            rejected_class_folder_path = os.path.join(rejected_folder_path, class_name,str(thresholds))
            os.makedirs(rejected_class_folder_path, exist_ok=True)

            
            rejected_filename = f"{image_filename.split('.')[0]}_{xmin}_{ymin}_{xmax}_{ymax}.jpg"
            rejected_path = os.path.join(rejected_class_folder_path, rejected_filename)

            
            cropped_image.save(rejected_path)

            print(f"Rejected image saved: {rejected_filename}")